/*
  Author: Chen Haodong (shinkunyachen@outlook.com)
  -----
  Copyright (c)2021 - present Chen Haodong. All Rights Reserved.
 */

#ifndef BRANCH_ANALYZER_H
#define BRANCH_ANALYZER_H

#include "LCFG.h"
#include "framework/BasicChecker.h"
#include <iostream>
#include <map>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>

enum ComparisonOp
{
    EQ,
    NEQ,
    GT,
    GE,
    LT,
    LE
};
class ConditionProbabilityTable
{
private:
    struct Condition
    {
        const clang::Expr *lhs;
        const clang::Expr *rhs;
        ComparisonOp op;
    };
    using Guesser = double (*)(const Condition &);

public:
    static const clang::Expr *remove_cast(const clang::Expr *expr);
    ConditionProbabilityTable();
    double probability(const clang::Expr *expr);

private:
    bool extract_cond(const clang::Expr *expr, Condition &cond);
    double guess(const Condition &cond);
    void init_guessers();
    //
    std::vector<Guesser> _guessers;
};
class BranchAnalyzer : public BasicChecker
{
public:
    BranchAnalyzer(ASTResource *resource, ASTManager *manager,
                   CallGraph *call_graph, Config *configure)
        : BasicChecker(resource, manager, call_graph, configure) {
        load_critical_path();
    }

    void analyze(ASTFile *file);

    const std::unordered_map<const clang::Stmt *, double> &results() const { return _results; }

private:
    int calc_basic_block_weight(const clang::FunctionDecl *caller, const LCFGBasicBlock *bb);
    void calc_basic_block_costs(const LCFG &lcfg, std::vector<int> &costs);
    double calc_cost_probability(const LCFGStatement *cond,
                                 const std::vector<int> &costs);
    double calc_cond_probability(const LCFGStatement *cond);
    double calc_cond_probability_impl(const clang::Expr *expr);
    template <typename _Pred>
    static bool contain_expr(const clang::Expr *expr, _Pred &&predicate);
    void report(clang::FunctionDecl *func_decl, const clang::Stmt *stmt,
                double probability);
    void load_critical_path();
    bool is_in_critical_path(const clang::FunctionDecl *caller, const clang::Expr *expr);
    bool is_in_critical_path(const clang::FunctionDecl *caller, const clang::Stmt *stmt);

private:
    std::unordered_map<const clang::Stmt *, double> _results;
    ConditionProbabilityTable _condition_table;
    std::unordered_multimap<std::string, std::string> _critical_paths;
};

#endif